"""
2019-3-1
增加生成含多个数字的图片功能，包含在函数random_multiple中，图片生成的代码由于不好重用之前
生成单个图片的函数，故生成图片的代码也塞到random_multiple函数里了。实际上，这部分代码相当
于是create_and_save_image()的改进版本。

若要调整图片中数字的个数，修改random_multiple的参数即可

2019-2-25
这个程序用来随机抽选北理工数据集中的部分测试集数据，生成图片，以便于复制到手机用于上传识别。
每个数字随机抽取5个。
"""
#! usr/bin/env python3
# coding=utf-8

from PIL import Image
import random
import os
import numpy as np
import re

DATA_DIR = os.path.join('dataSet', 'digits', 'testDigits')
IMG_DIR = os.path.join('BITDataeSet_images')
RANDOM_NUM = 5  # How many pictures each digit (single digit)
SCALE = 10  # scale ratio

PADDING = 5  # px, the value when scale=1
DIGIT_SIZE = 32, 32  # px, the width and height of a digit in the data set
IMG_SIZE = (DIGIT_SIZE[0] + PADDING) * SCALE, (DIGIT_SIZE[1] + PADDING) * SCALE  # the size of one image in the data set(single digit)


def choose_file(file_list):
    """Randomly choose the file name from the given file_list."""
    selected_list = [[]] * 10  # the list of the selected files
    num = 0
    for i in file_list:
        num += len(i)

    if num < RANDOM_NUM * 10:
        print("file_list might has not enough items, this might "
              "cause endless loop.")

    for i in range(10):
        n = 0
        while n < RANDOM_NUM:
            f_name = random.choice(file_list)
            if f_name[0] != str(i):
                continue
            elif f_name in selected_list[i]:
                continue
            else:
                selected_list[i].append(f_name)
                n += 1

    return selected_list


def load_data(path):
    """Load the data from the file, return the data as numpy ndarray.
        This implementation has some problem in efficiency.
    """
    img = np.zeros(IMG_SIZE, dtype=int)
    with open(path, 'r') as f:
        ln = 0  # which line is currently on
        for line in f:
            line = re.sub(r'[\n\r]', '', line)  # strip the \r or \n or \r\n
            for i in range(SCALE):
                col = 0  # which column is currently on
                for value in line:
                    for j in range(10):
                        img[ln][col] = value
                        col += 1
                ln += 1

    return img


def load_data_simple(path):
    img = np.zeros(DIGIT_SIZE, dtype=int)

    with open(path, 'r') as f:
        lines = f.readlines()
        for y in range(DIGIT_SIZE[1]):
            for x in range(DIGIT_SIZE[0]):
                img[y][x] = lines[y][x]

        return img


def create_and_save_image(f_name, img_arr):
    """Create a image using the givent img_arr,
     and save save it with the given file name, without suffix."""

    img = Image.fromarray(img_arr)
    # band = img.split()[0]
    img = img.convert('L')
    bands = []
    for i in range(4):  # 4 bands of RGBA mode
        bands.append(img.copy())
    bands[3] = bands[3].point(lambda x: x*255)  # set the alpha band to transparent

    img = Image.merge(mode='RGBA', bands=bands)
    img.save('%s.png' % f_name)

    return img


def random_single():
    """Randomly generate single picture with specific quantity."""
    f_list = os.listdir(DATA_DIR)  # the list of all the files
    f_list = choose_file(f_list)

    path_img = os.path.join(IMG_DIR, 'single')
    if not os.path.isdir(path_img):
        os.makedirs(path_img)

    print("Generating single digit pictures at path: %s..." % path_img)
    for i in f_list:
        for f_name in i:
            path_data = os.path.join(DATA_DIR, f_name)
            arr = load_data(path_data)

            file_name = f_name.split('.')[0]
            file_name = '%s.png' % file_name
            file_name = os.path.join(path_img, file_name)
            create_and_save_image(file_name, arr)

    print('Single all done !')


def random_multiple(num_of_digits):
    """Generate pictures of multiple digits. Randomly choose the digits from data set."""
    f_list = os.listdir(DATA_DIR)  # the list of all the files
    f_list = choose_file(f_list)
    flatten = (lambda lst: [x for i in lst for x in i])
    f_list = flatten(f_list)

    # choose digit files
    digit_f_list = []
    for i in range(num_of_digits):
        digit_f_list.append(random.choice(f_list))

    # load data set
    imgs_arr = []
    for i in range(num_of_digits):
        img = load_data_simple(os.path.join(DATA_DIR, '0_0.txt'))
        imgs_arr.append(img)

    # output image
    padding = PADDING * SCALE
    img_size = (2*PADDING + num_of_digits * DIGIT_SIZE[0]) * SCALE, (2*PADDING + DIGIT_SIZE[1]) * SCALE
    img = np.zeros((img_size[1], img_size[0]), dtype=int)
    for i in range(num_of_digits):
        path_data = os.path.join(DATA_DIR, digit_f_list[i])
        arr = load_data_simple(path_data)

        for y in range(DIGIT_SIZE[1]):
            for x in range(DIGIT_SIZE[0]):
                left = padding + i * SCALE * DIGIT_SIZE[0] + x * SCALE
                top = padding + y * SCALE
                right = padding + i * SCALE * DIGIT_SIZE[0] + (x+1) * SCALE
                bottom = padding + (y+1) * SCALE

                img[top:bottom, left:right] = arr[y][x]

    img = Image.fromarray(img)
    img = img.convert('L')
    img = img.point(lambda value: 255*value, mode='L')
    bands = []
    for i in range(4):
        bands.append(img.copy())

    return img


if __name__ == '__main__':
    np.set_printoptions(threshold=np.inf, linewidth=500)
    # random_single()

    path_img = os.path.join(IMG_DIR, 'multiple')
    if not os.path.isdir(path_img):
        os.makedirs(path_img)
    for i in range(10):
        im = random_multiple(5)
        im.save(os.path.join(path_img, '%i.png' % i))


